import torch
import torch.fft
import numpy as np
import tensorflow as tf
import time
import timeit 
import cupy as cp
import pyfftw
from tensorflow import fft
####### numpy tensorflow ###########################
def tf_cpu(n):
    tf.signal.fft3d(n)
def tf_gpu(n):
    tf.signal.fft3d(n)
def np_cpu(n):
    np.fft.fftn(n)

n1=np.load("YIG_n.npy")[:,:,:,1]
#=======================================
with tf.device('/cpu:0'):
    n1_cpu=tf.constant(n1,dtype=tf.complex64)
    tf_cpu_=np.array(timeit.repeat(stmt="tf_cpu(n1_cpu)", number=10,repeat=100, globals=globals())[1:])
with tf.device('/gpu:0'):
    n1_gpu=tf.constant(n1,dtype=tf.complex64)
    tf_gpu_=np.array(timeit.repeat(stmt="tf_gpu(n1_gpu)", number=10,repeat=100, globals=globals())[1:])
np_cpu_=np.array(timeit.repeat(stmt="np_cpu(n1)", number=10,repeat=100, globals=globals())[1:])
np.save("tf_cpu"+str(n1.shape)+".npy",tf_cpu_)
np.save("tf_gpu"+str(n1.shape)+".npy",tf_gpu_)
np.save("np_cpu"+str(n1.shape)+".npy",np_cpu_)
print("tfcpu",np.mean(tf_cpu_))
print("tfgpu",np.mean(tf_gpu_))
print("npcpu",np.mean(np_cpu_))

#====torch============================
def torch_(n):
    torch.fft.fftn(n)
def cupy_(n):
    cp.fft.fftn(n)
def fftw3_(n):
    fftwp1(n)
a = pyfftw.empty_aligned(n1.shape, dtype='complex128')
fftwp1 = pyfftw.builders._utils._FFTWWrapper(a,a,
    flags=["FFTW_EXHAUSTIVE"], threads=3
    )

n1_torchcpu=torch.tensor(n1)
n1_torchgpu=n1_torchcpu.cuda()
torch_cpu_=np.array(timeit.repeat(stmt="torch_(n1_torchcpu)", number=10,repeat=100, globals=globals())[1:])
torch_gpu_=np.array(timeit.repeat(stmt="torch_(n1_torchgpu)", number=10,repeat=100, globals=globals())[1:])
np.save("torch_cpu"+str(n1.shape)+".npy",torch_cpu_)
np.save("torch_gpu"+str(n1.shape)+".npy",torch_gpu_)
print("torchcpu",np.mean(torch_cpu_))
print("torchgpu",np.mean(torch_gpu_))
n1_cupy=cp.asarray(n1)
cp_gpu_=np.array(timeit.repeat(stmt="cupy_(n1_cupy)", number=10,repeat=100, globals=globals())[1:])
np.save("cp_gpu"+str(n1.shape)+".npy",cp_gpu_)
print("cpgpu",np.mean(cp_gpu_))
fftw3=np.array(timeit.repeat(stmt="fftw3_(n1)", number=10,repeat=100, globals=globals())[1:])
np.save("fftw3"+str(n1.shape)+".npy",fftw3)
print("fftw3",np.mean(fftw3))